import pytest
from fastapi.testclient import TestClient
from main import app
import uuid
from qdrant_client.http import models
from fastapi import FastAPI, UploadFile, File
from pymongo import MongoClient
from qdrant_client import QdrantClient
from qdrant_client.http import models
import requests
from typing import List, Dict, Any

@pytest.fixture
def client():
    return TestClient(app)

# MongoDB 客户端
mongo_client = MongoClient("mongodb://localhost:27017/")
db = mongo_client["rag_system"]
collection = db["documents"]

# Qdrant 客户端
qdrant_client = QdrantClient("http://localhost:6333")

# 确保Qdrant集合存在
def ensure_qdrant_collection():
    try:
        qdrant_client.get_collection("documents")
    except Exception:
        qdrant_client.create_collection(
            collection_name="documents",
            vectors_config=models.VectorParams(
                size=1024,  # 调整为实际的向量维度
                distance=models.Distance.COSINE
            )
        )

# 在应用启动时创建集合
ensure_qdrant_collection()

class OllamaClient:
    def __init__(self):
        self.base_url = "http://localhost:11434"

    def generate_vector(self, text: str) -> List[float]:
        response = requests.post(
            f"{self.base_url}/api/embeddings",
            json={
                "model": "bge-large",
                "prompt": text
            }
        )
        response.raise_for_status()
        return response.json()["embedding"]

    def generate_response(self, context: str, query: str) -> str:
        response = requests.post(
            f"{self.base_url}/api/generate",
            json={
                "model": "deepseek-r1:1.5b",
                "prompt": f"Context: {context}\n\nQuestion: {query}\n\nAnswer:",
                "stream": False
            }
        )
        response.raise_for_status()
        return response.json()["response"]

ollama_client = OllamaClient()

def upload_document_to_qdrant(doc_id: str, content: str, vector: List[float]):
    point_id = str(uuid.uuid4())  # 使用UUID作为Qdrant的点ID
    point = models.PointStruct(
        id=point_id,
        vector=vector,
        payload={"content": content}
    )
    qdrant_client.upsert(
        collection_name="documents",
        points=[point]
    )

@app.post("/upload/")
async def upload_document(file: UploadFile = File(...)):
    content = await file.read()
    content_str = content.decode('utf-8')
    
    # 存储到MongoDB
    doc_id = str(collection.insert_one({"content": content_str}).inserted_id)
    
    # 生成向量并存储到Qdrant
    vector = ollama_client.generate_vector(content_str)
    upload_document_to_qdrant(doc_id, content_str, vector)
    
    return {"message": "Document uploaded successfully"}

@app.get("/query/")
async def query_document(query: str):
    # 生成查询向量
    query_vector = ollama_client.generate_vector(query)
    
    # 从Qdrant检索相关文档
    search_results = qdrant_client.search(
        collection_name="documents",
        query_vector=query_vector,
        limit=3
    )
    
    # 构建上下文
    context = " ".join([hit.payload["content"] for hit in search_results])
    
    # 生成回答
    answer = ollama_client.generate_response(context, query)
    
    return {"answer": answer}